# Author : ZZH
# Date : 2025/5/23
# !/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :yolo_val_simple.py
# @Time      :2025/05/20
# @Author    :雨霓同学
# @Function  :YOLOv8 简易验证脚本，用于交通工具检测任务。支持基本验证功能，输出 mAP 指标，
#             供学生基于此迭代开发，添加日志、数据集信息、硬件信息等功能。

import argparse  # 命令行参数解析模块
from ultralytics import YOLO  # YOLOv8 模型接口


def parse_args():
    """
    解析命令行参数，配置验证所需的参数。

    Returns:
        argparse.Namespace: 包含所有命令行参数的对象
    """
    parser = argparse.ArgumentParser(description="YOLOv8 Simple Validation for Safety Helmet Detection")
    parser.add_argument('--data', type=str, default='../configs/data.yaml',
                        help='数据集配置文件路径（如 data.yaml，包含验证集路径和类别信息）')
    parser.add_argument('--weights', type=str, default='../models/checkpoints/train8-20250529-152642-yolov8m-best.pt',
                        help='训练好的模型权重文件（如 best.pt）')
    parser.add_argument('--imgsz', type=int, default=640,
                        help='输入图像尺寸，影响模型输入（默认: 640）')
    parser.add_argument('--device', type=str, default='cpu',
                        help='计算设备，如 0（GPU）或 cpu（默认: 0）')
    parser.add_argument('--run-name', type=str, default='val',
                        help='验证运行名称，影响输出目录名（默认: val）')
    return parser.parse_args()


def main(args):
    """
    主验证函数，执行模型加载和验证。

    Args:
        args (argparse.Namespace): 命令行参数对象，包含验证配置

    Returns:
        None

    Raises:
        FileNotFoundError: 如果 data.yaml 或模型文件不存在
        Exception: 其他验证错误
    """
    # 加载 YOLOv8 模型
    # 使用 ultralytics 的 YOLO 类，传入训练好的模型权重（如 best.pt）
    model = YOLO(args.weights)

    # 执行验证
    # model.val 是 YOLOv8 的验证接口，评估模型在验证集上的性能
    # 数据集路径（data）来自 data.yaml，指定验证集
    # imgsz 设置输入图像尺寸，device 指定 GPU 或 CPU
    # project 和 name 控制输出目录（如 runs/val/val）
    results = model.val(
        data=args.data,  # 数据集配置文件路径
        imgsz=args.imgsz,  # 输入图像尺寸
        device=args.device,  # 计算设备
        project='runs/val',  # 输出目录（runs/val/）
        name=args.run_name  # 运行名称（子目录，如 val）
    )

    # 输出验证结果
    # results.box.map50 是 mAP@50 指标，results.box.map 是 mAP@50:95 指标
    print(f"mAP@50: {results.box.map50:.3f}")
    print(f"mAP@50:95: {results.box.map:.3f}")


if __name__ == '__main__':
    # 解析命令行参数
    args = parse_args()
    # 调用主函数，开始验证
    main(args)
